# -*- coding: utf-8 -*-
# time: 2025/4/9 14:52
# file: llm_ts.py
# author: hanson
"""

测试微调后的模型
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载微调后的模型和tokenizer
model_path = "./qwen2-0.5b-finetuned"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
)

# 如果微调时使用了LoRA，推理时需要合并适配器或显式加载Peft模型
model = model.merge_and_unload()
# 构造测试提示
prompt = """Instruction: 用简单的中文解释人工智能
Input: 面向小学生
Output:"""

# 生成回答
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    do_sample=True,
    temperature=0.7,
    top_p=0.9
)

# 解码输出
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("模型回答：\n", response.split("Output:")[-1].strip())